package com.vr.xxaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.vr.xxaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 用于处理具有增强抽象的工具/函数调用的基代理类
 *
 * @author hzh
 * @date 2025/06/08
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    /** 可用工具 */
    private final ToolCallback[] availableTools;
    /** 保存工具响应信息 */
    private ChatResponse toolCallchatResponse;
    /** 工具调用管理器,这样就可以介入工具的调用 */
    private final ToolCallingManager toolCallingManager;
    /** 大模型调用的配置，可以关闭工具自动调用的机制，改为手动,自己维护上下文 */
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        // 工具尤其创建者停供
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                // 禁用spring ai内置的工具调用机制,交给自己维护选项和消息上下文
                .withProxyToolCalls(true)
                .build();

    }

    /**
     * 使用工具处理当前状态并决定下一步行动
     *
     * @return boolean
     */
    @Override
    public boolean think() {
        if (StrUtil.isNotEmpty(this.getNextStepPrompt())) {
            this.getMessageList().add(new UserMessage(this.getNextStepPrompt()));
        }
        List<Message> messageList = this.getMessageList();
        Prompt prompt = new Prompt(messageList, this.chatOptions);
        try {
            // 根据聊天的返回值，获取工具选项的调用
            ChatResponse chatResponse = this.getLlm().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(this.availableTools)
                    .call()
                    .chatResponse();
            // 记录响应，用于act
            this.toolCallchatResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            // 输出提示日志
            String result = assistantMessage.getText();
            // 调用的工具列表
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            log.info("{}思考了：{}", this.getName(), result);
            log.info("{}选择了：{}个工具来使用", this.getName(), toolCalls.size());
            String toolCallInfo = toolCalls.stream().map(toolCall -> String.format("工具名称：%s，参数：%s",
                    toolCall.name(),
                    toolCall.arguments()
            )).collect(Collectors.joining("\n"));
            log.info(toolCallInfo);
            if (toolCalls.isEmpty()) {
                // 没有工具调用时，手动记录助手消息,不需要act
                getMessageList().add(assistantMessage);
                return false;
            } else {
                // 有工具调用时，会自动记录助手消息，不需要手动记录，需要act
                // AssistantMessage assistantMessage = toolCallGeneration.get().getOutput();
                return true;
            }
        } catch (Exception e) {
            log.error("{}思考中出现的问题：{}", this.getName(), e.getMessage());
            // 记录错误信息
            this.getMessageList().add(new AssistantMessage(String.format("处理时遇到了错误:%s", e.getMessage())));
            return false;
        }
    }

    /**
     * 执行工具调用并执行
     *
     * @return {@code String }
     */
    @Override
    public String act() {
        if (!this.toolCallchatResponse.hasToolCalls()) {
            return "没有工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(this.getMessageList(), this.chatOptions);
        ToolExecutionResult toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, this.toolCallchatResponse);
        // 记录上下文消息,conversationHistory已经包含了助手和工具的返回结果
        this.getMessageList().addAll(toolExecutionResult.conversationHistory());
        //List<Message> conversationHistory = buildConversationHistoryAfterToolExecution(prompt.getInstructions(),
        // 				assistantMessage, internalToolExecutionResult.toolResponseMessage());
        // 通过源码可知，最后一个元素是工具的返回结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> String.format("工具：%s,完成了任务,结果：%s", response.name(), response.responseData()))
                .collect(Collectors.joining("\n"));
        // 判断一下是否调用了终止工具
        boolean doTerminate = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> response.name().equals("doTerminate"));
        if (doTerminate) {
            // 表示调用了终止工具，手动完成
            this.setState(AgentState.FINISHED);
        }
        log.info(results);
        return results;
    }
}
